import socket
import sys
import argparse
import threading
import time
import random

class ChatClient:
    def __init__(self, ip, port):
        self.__id = None
        self.__remote_addr = None
        self.__server_addr = (ip, port)
        self.client_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

    def keep_alive(self):
        while self.__id == None:
            # req_id_c2s
            self.client_socket.sendto(bytes([0x00]), self.__server_addr)
            time.sleep(1)
        heartbeat_packet = bytearray([0x01]) + self.__id.to_bytes(4, "little")
        while True:
            # keep_alive_c2s
            self.client_socket.sendto(heartbeat_packet, self.__server_addr)
            time.sleep(60)

    def sender(self):
        while True:
            mode = input("Please input mode(0: connect, 1: chat):")
            mode = int(mode)
            match mode:
                case 0:
                    remote_id = input("Please input remote id:")
                    remote_id = int(remote_id)
                    self.client_socket.sendto(bytearray([0x02]) + remote_id.to_bytes(4, "big"), self.__server_addr)
                case 1:
                    if self.__remote_addr == None:
                        print("no remote addr.")
                        continue
                    q_flag = 0
                    while q_flag == 0:
                        str = input()
                        if str == "q":
                            q_flag = 1
                            break
                        self.client_socket.sendto(str.encode("utf-8"), self.__remote_addr)

    def worker(self, data, addr):
        if addr == self.__server_addr:
            print("data from server.")
            code_b = data[0:1]
            code = int.from_bytes(code_b, "little")
            buf = data[1:]
            match code:
                case 0x00:
                    # id_s2c
                    self.__id = int.from_bytes(buf, "little")
                    print("get id: %d" % self.__id)
                case 0x01:
                    # update_remote_addr_s2c
                    addr_list = buf.decode("utf-8").split(":")
                    self.__remote_addr = (addr_list[0], int(addr_list[1]))
                    print("update remote addr: %s:%d" % (self.__remote_addr[0], self.__remote_addr[1]))
        else:
            print("data from other client: %s" % (data.decode("utf-8")))

    def start(self):
        print("Start chat client which server addr is {self.__server_addr[0]}:{self.__server_addr[1]}")
        self.client_socket.setblocking(True)
        t_keeper = threading.Thread(target=self.keep_alive, name="keeper")
        t_keeper.start()
        t_sender = threading.Thread(target=self.sender, name="sender")
        t_sender.start()
        while True:
            data, addr = self.client_socket.recvfrom(1024)
            threading.Thread(target=self.worker, args=(data, addr)).start()

class ChatServer:
    def __init__(self, host, port):
        self.clients = {}
        self.host = host
        self.port = port
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

    class Client:
        def __init__(self, addr):
            self.addr = addr
            self.partner = None

        def set_partner(self, partner):
            self.partner = partner

        def get_partner(self):
            return self.partner

        def has_partner(self):
            return self.partner != None

        def get_addr(self):
            return self.addr

        def set_addr(self, addr):
            self.addr = addr

    def worker(self, data, addr):
        print(f"Received message from {addr}")
        code_b = data[0:1]
        code = int.from_bytes(code_b, "little")
        buf = data[1:]
        match code:
            case 0x00:
                # req_id_c2s
                id = random.randint(1, 1000000000)
                self.clients[id] = self.Client(addr)
                print("new client: %d: (%s:%d)" % (id, addr[0], addr[1]))
                self.server_socket.sendto(bytes([0x00]) + id.to_bytes(4, "little"), addr)
            case 0x01:
                # keep_alive_c2s
                id = int.from_bytes(buf, "little")
                if self.clients[id].get_addr() != addr:
                    # update remote addr
                    print("client %d addr update: (%s:%d)." % (id, addr[0], addr[1]))
                    self.clients[id].set_addr(addr)
                    if self.clients[id].has_partner():
                        self.server_socket.sendto(bytes([0x01]) + (self.clients[id].get_partner().get_addr()[0] + ":" + str(self.clients[id].get_partner().get_addr()[1])).encode("utf-8"), addr)
            case 0x02:
                # connect_c2s
                id = int.from_bytes(buf, "little")
                print("connect request to %d" % id)
                print(self.clients)
                if id in self.clients:
                    client_a_id = None
                    client_a = None
                    client_b = self.clients[id]
                    for (key, value) in self.clients.items():
                        if value.get_addr() == addr:
                            client_a_id = key
                            client_a = value
                            break
                    client_a.set_partner(client_b)
                    client_b.set_partner(client_a)
                    print("connect %d and %d" % (client_a_id, id))
                    self.server_socket.sendto(bytes([0x01]) + (client_b.get_addr()[0] + ":" + str(client_b.get_addr()[1])).encode("utf-8"), client_a.get_addr())

    def start(self):
        self.server_socket.bind((self.host, self.port))
        self.server_socket.setblocking(True)
        print(f"Server bind on {self.host}:{self.port}")
        while True:
            data, addr = self.server_socket.recvfrom(1024)
            threading.Thread(target=self.worker, args=(data, addr)).start()

def start_server(host, port):
    server = ChatServer(host, port)
    server.start()

def start_client():
    client = ChatClient("101.200.162.149", 25565)
    client.start()

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--mode", choices=["server", "client"], help="Choose mode: server or client", default="client")
    parser.add_argument("--host", type=str, help="Host address", default="0.0.0.0")
    parser.add_argument("--port", type=int, help="Port number", default=25565)
    args = parser.parse_args()
    mode = args.mode
    host = args.host
    port = args.port

    print("Starting chat %s with host: %s, port: %d" % (mode, host, port))

    if mode == "server":
        start_server(host, port)
    elif mode == "client":
        start_client()
    else:
        print("Invalid mode")
        sys.exit(1)

if __name__ == "__main__":
    main()
